/** Utility for building query timeout call options. */
public final class QueryTimeoutCallOption implements CallOption {
  private QueryTimeoutCallOption() {}

    /**
     * Merges base call options with extra options.
     *
     * @param baseOptions the base options
     * @param extraOptions additional options to merge
     * @return combined array of options
     */
    public static CallOption[] mergeOptions(CallOption[] baseOptions, CallOption... extraOptions) {
        if (extraOptions.length == 0) {
            return baseOptions;
        }
        final CallOption[] combined = new CallOption[baseOptions.length + extraOptions.length];
        System.arraycopy(baseOptions, 0, combined, 0, baseOptions.length);
        System.arraycopy(extraOptions, 0, combined, baseOptions.length, extraOptions.length);
        return combined;
    }

  /**
   * Builds a CallOption array from a statement's timeout configuration.
   *
   * @param statement the statement to extract timeouts from
   * @return CallOption array with timeout header, or empty array if no timeout configured
   * @throws SQLException on error retrieving timeout values
   */
  public static CallOption[] fromStatement(ArrowFlightInfoStatement statement)
      throws SQLException {
    final Duration statementTimeout = Duration.ofSeconds(statement.getQueryTimeout());
    final Duration connectionTimeout =
        Duration.ofSeconds(statement.getConnection().getConnectionQueryTimeoutSeconds());
    return build(statementTimeout, connectionTimeout);
  }

  /**
   * Builds a CallOption array from a statement ID and connection.
   *
   * @param statementId the statement ID to look up in the connection's statement map
   * @param connection the Avatica connection to extract timeouts from
   * @return CallOption array with timeout header, or empty array if no timeout configured
   */
  public static CallOption[] fromStatementHandle(int statementId, AvaticaConnection connection) {
    Duration statementTimeout = Duration.ZERO;
    final AvaticaStatement stmt = connection.statementMap.get(statementId);
    if (stmt != null) {
      try {
        statementTimeout = Duration.ofSeconds(stmt.getQueryTimeout());
      } catch (SQLException ignored) {
        // fall through to connection-level timeout
      }
    }
    Duration connectionTimeout = Duration.ZERO;
    if (connection instanceof ArrowFlightConnection) {
      connectionTimeout =
          Duration.ofSeconds(
              ((ArrowFlightConnection) connection).getConnectionQueryTimeoutSeconds());
    }
    return build(statementTimeout, connectionTimeout);
  }

  /**
   * Builds a CallOption array containing query timeout header if a timeout is configured.
   *
   * @param statementTimeout timeout at statement level (ZERO means use connection-level)
   * @param connectionTimeout timeout at connection level
   * @return CallOption array with timeout header, or empty array if no timeout configured
   */
  private static CallOption[] build(Duration statementTimeout, Duration connectionTimeout) {
    final Duration effectiveTimeout =
        !statementTimeout.isZero() ? statementTimeout : connectionTimeout;

      if (effectiveTimeout.isZero()) {
          return new CallOption[0]; // short-circuit
      }

      final FlightCallHeaders headers = new FlightCallHeaders();
      headers.insert("x-query-timeout-ms", String.valueOf(effectiveTimeout.toMillis()));
      return new CallOption[] {new HeaderCallOption(headers)};
  }
}